import librosa
import soundfile as sf
import numpy as np
import logging
import random
from scipy.signal import butter, lfilter
from config import config
from utils import setup_logger

logger = setup_logger()



class AudioProcessor:
    def __init__(self):
        self.sample_rate = config.sample_rate  # 从配置中读取采样率

    def butter_lowpass(self, cutoff, fs, order=None):
        """创建低通滤波器"""
        if order is None:
            order = config.lowpass_order  # 使用配置的滤波器阶数
        nyq = 0.5 * fs
        normal_cutoff = cutoff / nyq
        b, a = butter(order, normal_cutoff, btype='low', analog=False)
        return b, a

    def apply_lowpass(self, data, cutoff, fs, order=None):
        """应用低通滤波器"""
        b, a = self.butter_lowpass(cutoff, fs, order=order)
        y = lfilter(b, a, data)
        return y

    def process_audio(self, input_path, output_path):
        """处理音频：音高变换、速度调整、添加噪声"""
        try:
            # 加载音频
            y, sr = librosa.load(input_path, sr=self.sample_rate)
            logger.info(f"加载音频: {input_path}, 长度: {len(y)/sr:.2f}s")

            # 获取配置参数
            pitch_factor = random.uniform(*config.pitch_factor_range)
            speed_factor = random.uniform(*config.speed_factor_range)
            noise_intensity = float(config.noise_intensity)  # 确保是浮点数
            mask_width_ms = int(config.mask_width_ms)  # 确保是整数
            mask_interval_s = float(config.mask_interval_s)  # 确保是浮点数
            mask_intensity_min, mask_intensity_max = config.mask_intensity_range
            lowpass_cutoff_min, lowpass_cutoff_max = config.lowpass_cutoff_range
            lowpass_order = int(config.lowpass_order)  # 确保是整数

            # 音高变换
            y_pitch = librosa.effects.pitch_shift(y, sr=sr, n_steps=pitch_factor)

            # 速度调整
            y_speed = librosa.effects.time_stretch(y_pitch, rate=speed_factor)

            # 添加噪声 - 降低噪声强度
            noise_intensity = float(config.noise_intensity) * 0.3  # 降低70%的噪声强度
            noise = np.random.normal(0, noise_intensity, y_speed.shape)
            y_final = y_speed + noise

            # 音频掩码
            mask_width_samples = int(mask_width_ms * sr / 1000)
            mask_interval_samples = int(mask_interval_s * sr)
            mask_intensity = random.uniform(mask_intensity_min, mask_intensity_max)

            for i in range(0, len(y_final), mask_interval_samples):
                end_idx = min(i + mask_width_samples, len(y_final))
                y_final[i:end_idx] *= mask_intensity

            # 低通滤波
            cutoff = random.uniform(lowpass_cutoff_min, lowpass_cutoff_max)
            y_final = self.apply_lowpass(y_final, cutoff, sr, order=lowpass_order)

            # 音量标准化 - 防止音量过小
            # 计算音频的RMS值
            rms = np.sqrt(np.mean(y_final**2))
            if rms > 0:
                # 将音量调整到原始音量的80%左右，避免过度衰减
                target_rms = np.sqrt(np.mean(y**2)) * 0.8
                gain_factor = target_rms / rms

                # 限制增益范围，避免过度放大
                gain_factor = np.clip(gain_factor, 0.5, 2.0)
                y_final = y_final * gain_factor

            # 防止削波失真
            max_val = np.max(np.abs(y_final))
            if max_val > 0.95:
                y_final = y_final * (0.95 / max_val)

            # 确保输出文件扩展名为.wav
            if not output_path.endswith('.wav'):
                output_path = output_path.replace('.mp4', '.wav')

            # 保存音频
            sf.write(output_path, y_final, int(sr), format='WAV')
            logger.info(f"音频处理完成: {output_path}")

        except Exception as e:
            logger.error(f"音频处理失败: {str(e)}")
            raise
